{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%reload_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "import seaborn as sns\n",
    "\n",
    "import sys\n",
    "sys.path.append(os.path.abspath('../'))\n",
    "from common import *\n",
    "import configs\n",
    "import opt_configs\n",
    "\n",
    "from constants import OUTPUT_DIR\n",
    "import matplotlib.gridspec as gridspec\n",
    "\n",
    "def gallery(array, ncols=2):\n",
    "    \"\"\"Code from https://stackoverflow.com/a/42041135/2351867\"\"\"\n",
    "    nindex, height, width, intensity = array.shape\n",
    "    nrows = nindex // ncols\n",
    "    assert nindex == nrows*ncols\n",
    "    result = (array.reshape(nrows, ncols, height, width, intensity)\n",
    "              .swapaxes(1, 2)\n",
    "              .reshape(height*nrows, width*ncols, intensity))\n",
    "    return result\n",
    "\n",
    "def set_spine_color(ax, color='white'):\n",
    "    ax.spines['bottom'].set_color(color)\n",
    "    ax.spines['top'].set_color(color) \n",
    "    ax.spines['right'].set_color(color)\n",
    "    ax.spines['left'].set_color(color)\n",
    "\n",
    "\n",
    "darkmode = False\n",
    "write_png = False\n",
    "\n",
    "linecolor = 'white' if darkmode else 'black'\n",
    "\n",
    "sns.set()\n",
    "sns.set_style('darkgrid', {'axes.edgecolor': 'black'})\n",
    "matplotlib.rcParams['axes.linewidth'] = 0.5\n",
    "\n",
    "fig_name = 'teaser'\n",
    "fig_dir = join(FIGURE_DIR, fig_name)\n",
    "used_configs = [configs.Warp()]\n",
    "names = ['Ours']\n",
    "\n",
    "scenes = []\n",
    "used_opt_configs = []\n",
    "names_pretty = []\n",
    "insets = []\n",
    "\n",
    "scenes.append('bench')\n",
    "used_opt_configs.append('diffuse-32-hq')\n",
    "\n",
    "show_iters = ['final']\n",
    "ref_views = [0, 4, 8, 10]\n",
    "\n",
    "exposure = 0.5\n",
    "fontsize = 16\n",
    "base_size = 2\n",
    "n_rows = len(scenes)\n",
    "n_cols = len(show_iters) + 3\n",
    "total_width = base_size * n_cols\n",
    "aspect = 0.98 * n_rows / n_cols\n",
    "aspect = 0.42\n",
    "total_width = TEXT_WIDTH\n",
    "fig = plt.figure(1, figsize=(total_width, aspect * total_width), constrained_layout=False)\n",
    "iter_rename_dict = {show_iters[-1]: 'Final', 0: 'Initialization'}\n",
    "\n",
    "gs = fig.add_gridspec(2, 1, wspace=0.0, hspace=0.0, height_ratios=[2, 1])\n",
    "\n",
    "gs0 = gridspec.GridSpecFromSubplotSpec(1, len(show_iters) + 3, subplot_spec=gs[0], wspace=0.025, hspace=0.0)\n",
    "inset_size = 0.2\n",
    "y_offset = -0.05\n",
    "crop_pixels = 40\n",
    "ranges = [5e-7, 2e-6, 5e-5]\n",
    "scene = scenes[0]\n",
    "opt_name = used_opt_configs[0]\n",
    "\n",
    "opt_config = opt_configs.get_opt_config(opt_name)\n",
    "output_dir = join(OUTPUT_DIR, scene, opt_config.name, used_configs[0].name) \n",
    "\n",
    "show_ref_as_subplots = False\n",
    "if show_ref_as_subplots:\n",
    "    gs_insets = gridspec.GridSpecFromSubplotSpec(2, 2, subplot_spec=gs0[0], wspace=0.0, hspace=0.0)\n",
    "all_ref_images = []\n",
    "for idx, ref_view in enumerate(ref_views):\n",
    "    r = idx // 2\n",
    "    c = idx % 2\n",
    "    img = read_img(join(output_dir, f'ref-{ref_view:02d}.exr'), exposure=exposure)\n",
    "    if all_ref_images is None:\n",
    "        all_ref_images = np.zeros((2 * img.shape[0], 2 * img.shape[1], 3))\n",
    "    all_ref_images.append(img)\n",
    "\n",
    "    if show_ref_as_subplots:\n",
    "        ax = fig.add_subplot(gs_insets[r, c])\n",
    "        ax.imshow(img, interpolation='none')\n",
    "        disable_ticks(ax)\n",
    "\n",
    "ax = fig.add_subplot(gs0[0])\n",
    "txt = ax.set_title('(a) Reference images', color=linecolor, fontsize=fontsize, y=y_offset, va='top')\n",
    "\n",
    "if not show_ref_as_subplots:\n",
    "    all_ref_images = gallery(np.stack(all_ref_images, axis=0))\n",
    "    all_ref_images[all_ref_images.shape[0]//2, :, :] = 1\n",
    "    all_ref_images[:, all_ref_images.shape[1]//2, :] = 1\n",
    "    ax.imshow(all_ref_images)\n",
    "    disable_ticks(ax)\n",
    "    set_spine_color(ax, color=linecolor)\n",
    "else:\n",
    "    ax.set_axis_off()\n",
    "\n",
    "col = 1\n",
    "r = None\n",
    "for cfg in used_configs:\n",
    "    ax = fig.add_subplot(gs0[col])\n",
    "    fn = join(FIGURE_DIR, fig_name, scene, f'warp_final_geo.exr')\n",
    "    img = read_img(fn, exposure=exposure)\n",
    "    img = img[crop_pixels:-crop_pixels, crop_pixels:-crop_pixels, :]\n",
    "    ax.imshow(img, interpolation='none', extent=[0, 1, 0, 1])\n",
    "    disable_ticks(ax)\n",
    "    txt = ax.set_title('(b) Optimized surface', color=linecolor, fontsize=fontsize, y=y_offset, va='top')\n",
    "    col += 1\n",
    "    set_spine_color(ax, color=linecolor)\n",
    "\n",
    "    ax = fig.add_subplot(gs0[col])\n",
    "    fn = join(FIGURE_DIR, fig_name, scene, f'{cfg.name}_final.exr')\n",
    "    img = read_img(fn, exposure=exposure)\n",
    "    img = img[crop_pixels:-crop_pixels, crop_pixels:-crop_pixels, :]\n",
    "    ax.imshow(img, interpolation='none', extent=[0, 1, 0, 1])\n",
    "    disable_ticks(ax)\n",
    "    txt = ax.set_title('(c) Optimization result', color=linecolor, fontsize=fontsize, y=y_offset, va='top')\n",
    "    col += 1\n",
    "    set_spine_color(ax, color=linecolor)\n",
    "\n",
    "    ax = fig.add_subplot(gs0[col])\n",
    "    fn = join(FIGURE_DIR, fig_name, scene, f'reference.exr')\n",
    "    img = read_img(fn, exposure=exposure)\n",
    "    img = img[crop_pixels:-crop_pixels, crop_pixels:-crop_pixels, :]\n",
    "    ax.imshow(img, interpolation='none', extent=[0, 1, 0, 1])\n",
    "    disable_ticks(ax)\n",
    "    txt = ax.set_title('(d) Ground truth', color=linecolor, fontsize=fontsize, y=y_offset, va='top')\n",
    "    col += 1\n",
    "    set_spine_color(ax, color=linecolor)\n",
    "\n",
    "# Add optimization states\n",
    "if True:\n",
    "    show_iters = np.arange(0, 512, 64).tolist() + [511,]\n",
    "    gs1 = gridspec.GridSpecFromSubplotSpec(1, len(show_iters), subplot_spec=gs[1], wspace=0.025, hspace=0.0)\n",
    "    scene = scenes[0]\n",
    "    opt_name = used_opt_configs[0]\n",
    "\n",
    "    opt_config = opt_configs.get_opt_config(opt_name)\n",
    "    output_dir = join(OUTPUT_DIR, scene, opt_config.name, used_configs[0].name)\n",
    "    cfg = configs.Warp()\n",
    "    r = None\n",
    "\n",
    "    iter_rename_dict = {show_iters[-1]: '512', 0: 'Initialization'}\n",
    "    for col, show_iter in enumerate(show_iters):\n",
    "        ax = fig.add_subplot(gs1[col])\n",
    "        suffix = f'{show_iter:04d}' if isinstance(show_iter, int) else show_iter\n",
    "\n",
    "        fn = join(FIGURE_DIR, fig_name, scene, f'{cfg.name}_{suffix}.exr')\n",
    "        img = read_img(fn, exposure=exposure)\n",
    "        topc = 40\n",
    "        bottomc = 2 * crop_pixels - topc\n",
    "        img = img[topc:-bottomc, crop_pixels:-crop_pixels, :]\n",
    "        ax.imshow(img, interpolation='none', extent=[0, 1, 0, 1])\n",
    "        disable_ticks(ax)\n",
    "        set_spine_color(ax, color=linecolor)\n",
    "\n",
    "        label = iter_rename_dict[show_iter] if show_iter in iter_rename_dict.keys() else f'{show_iter}'\n",
    "        txt = ax.set_title(label, y=0.9, x=0.025, color='k', ha='left', va='top', fontsize=14)\n",
    "        txt.set_path_effects([path_effects.withStroke(linewidth=0.5, foreground='white')])\n",
    "plt.margins(0, 0)\n",
    "\n",
    "if write_png:\n",
    "    kwargs = {'file_format': 'png', 'transparent': True, 'dpi': 250}\n",
    "\n",
    "# save_fig(fig_name, **kwargs)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10.4 ('mi')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.4"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "c8cc192903caa17681aa39d71202092ac11a526d37a1c4ad2948f13605924304"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
